import torch
import torch.nn.functional as F
from torch_geometric.nn import MessagePassing
from torch_geometric.data import Data

class EdgeUpdateLayer(torch.nn.Module):
    def __init__(self, model_dim, edge_channels):
        super(EdgeUpdateLayer, self).__init__()
        self.edge_mlp = torch.nn.Sequential(
            torch.nn.Linear(2 * model_dim + edge_channels, edge_channels),
            torch.nn.ReLU(),
            torch.nn.Linear(edge_channels, edge_channels)
        )

    def forward(self, x, edge_index, edge_attr):
        row, col = edge_index
        edge_inputs = torch.cat([x[:,row,:], x[:,col,:], edge_attr], dim=-1)
        edge_attr = self.edge_mlp(edge_inputs)
        return edge_attr